﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Security.Cryptography;
using System.Text;
using System.Text.RegularExpressions;

namespace NWebSockets
{
	public class WebSocketHandshakeException : Exception
	{
		/// <summary>
		/// Default constructor
		/// </summary>
		public WebSocketHandshakeException() : base( "The web socket handshake was invalid" ) { }
	}

	internal class WebSocketServerHandshake
	{
		#region Constructors
		/// <summary>
		/// Default constructor
		/// </summary>
		internal WebSocketServerHandshake() { }

		/// <summary>
		/// Construct a new server handshake from the incoming client handshake
		/// </summary>
		internal WebSocketServerHandshake(WebSocketClientHandshake clientHandshake)
		{
			// generate the properties from the client handshake
			this.Origin = clientHandshake.Origin;
			this.Location = string.Format( "ws://{0}{1}", clientHandshake.Host, clientHandshake.ResourcePath );
			this.Challenge = this.GetChallengeResponse( clientHandshake.Key1, clientHandshake.Key2, clientHandshake.Challenge );
		}
		#endregion

		#region Internal Methods
		/// <summary>
		/// Convert the handshake to a byte array for sending over the pipe
		/// </summary>
		internal byte[] ToBytes()
		{
			// construct the server handshake
			var handshake = string.Format( this.ServerHandshakeFormat, this.Origin, this.Location );
			var handshakeBytes = Encoding.UTF8.GetBytes( handshake ).ToList();

			// add on the challenge response
			handshakeBytes.AddRange( this.Challenge );

			return handshakeBytes.ToArray();
		}

		/// <summary>
		/// Equality comparison
		/// </summary>
		public override bool Equals( object obj )
		{
			// can't compare to null object
			if ( obj == null )
				return false;

			// try to convert the object to a handshake
			var handshake = obj as WebSocketServerHandshake;
			if ( handshake == null )
				return false;

			// otherwise, check to see if all the properties are equal
			return this.Origin == handshake.Origin &&
				   this.Location == handshake.Location &&
				   this.Challenge == handshake.Challenge;
		}

		/// <summary>
		/// Equality comparison for WebSocketServerHandshake type, for performance
		/// </summary>
		public bool Equals( WebSocketServerHandshake handshake )
		{
			// can't compare to null object
			if ( handshake == null )
				return false;

			// otherwise, check to see if all the properties are equal
			return this.Origin == handshake.Origin &&
				   this.Location == handshake.Location &&
				   this.Challenge == handshake.Challenge;
		}

		/// <summary>
		/// Overload for the hashcode
		/// </summary>
		public override int GetHashCode()
		{
			// get the hashcodes for all the properties
			return this.Origin.GetHashCode() +
				   this.Location.GetHashCode() +
				   this.Challenge.GetHashCode();
		}
		#endregion

		#region Internal Properties
		/// <summary>
		/// The origin of the web socket connection
		/// </summary>
		internal string Origin { get; set; }

		/// <summary>
		/// The location of server
		/// </summary>
		internal string Location { get; set; }

		/// <summary>
		/// A byte array generated from the client keys and challenge
		/// </summary>
		internal byte[] Challenge { get; set; }
		#endregion

		#region Private Methods
		/// <summary>
		/// Gets the response bytes for a particular request key
		/// </summary>
		private byte[] GetKeyResponseBytes( string key )
		{
			var keyDigits = key.Where( x => Char.IsDigit( x ) ).ToArray();

			var digitNumber = Int64.Parse( new String( keyDigits ) );
			var numberOfSpaces = key.Count( x => x == ' ' );

			var responseNumber = (Int32)( digitNumber / numberOfSpaces );
			var responseBytes = BitConverter.GetBytes( responseNumber );

			if ( BitConverter.IsLittleEndian )
				Array.Reverse( responseBytes );

			return responseBytes;
		}

		/// <summary>
		/// Generates the server handshake challenge (the last 16 bytes)
		/// </summary>
		/// <remarks>http://tools.ietf.org/html/draft-ietf-hybi-thewebsocketprotocol-03#section-1.3</remarks>
		private byte[] GetChallengeResponse( string key1, string key2, byte[] challenge )
		{
			var response = new List<byte>();
			response.AddRange( this.GetKeyResponseBytes( key1 ) );
			response.AddRange( this.GetKeyResponseBytes( key2 ) );
			response.AddRange( challenge );

			var hashedResponse = MD5.Create().ComputeHash( response.ToArray() );

			return hashedResponse;
		}
		#endregion

		/// <summary>
		/// A format string for the server handshake
		/// </summary>
		internal string ServerHandshakeFormat = "HTTP/1.1 101 WebSocket Protocol Handshake\r\n" + 
												"Upgrade: WebSocket\r\n" + 
												"Connection: Upgrade\r\n" + 
												"Sec-WebSocket-Origin: {0}\r\n" + 
												"Sec-WebSocket-Location: {1}\r\n" + 
												"\r\n";

		/// <summary>
		/// Generates a server handshake from the incoming client handshake
		/// </summary>
		internal static WebSocketServerHandshake FromClientHandshake( byte[] clientHandshakeBytes, int clientHandshakeSize )
		{
			// extract the client handshake from the buffer
			var clientHandshake = WebSocketClientHandshake.Parse( clientHandshakeBytes, clientHandshakeSize );

			// generate the server handshake
			var serverHandshake = new WebSocketServerHandshake(clientHandshake);

			return serverHandshake;
		}

		/// <summary>
		/// Parses out the server handshake from a buffer
		/// </summary>
		internal static WebSocketServerHandshake Parse( byte[] serverHandshakeBytes, int serverHandshakeSize )
		{
			// convert the handshake to text for regular expression parsing
			var handshake = Encoding.UTF8.GetString( serverHandshakeBytes );

			// construct the server handshake from the values found in the buffer
			var serverHandshake = new WebSocketServerHandshake()
			{
				Location = Regex.Match( handshake, @"Sec-WebSocket-Location: (?<location>.*)\r\n" ).Groups["location"].Value,
				Origin = Regex.Match( handshake, @"Origin: (?<origin>.*)\r\n" ).Groups["origin"].Value,
				Challenge = serverHandshakeBytes.Skip( serverHandshakeSize - 16 ).Take( 16 ).ToArray()
			};

			return serverHandshake;
		}
	}

	internal class WebSocketClientHandshake
	{
		#region Constructors
		/// <summary>
		/// Constructs a new client handshake
		/// </summary>
		internal WebSocketClientHandshake() 
		{
			this.Key1 = this.GenerateHandshakeKey();
			this.Key2 = this.GenerateHandshakeKey();
			this.Challenge = this.GetRandomBytes( 8 );
		}

		/// <summary>
		/// Constructs a handshake to a particular server resource
		/// </summary>
		internal WebSocketClientHandshake( string hostName, int port, string resourcePath ) : this()
		{
			this.Host = string.Format( "{0}:{1}", hostName, port );
			this.ResourcePath = resourcePath;
		}
		#endregion

		#region Internal Methods
		/// <summary>
		/// Converts the full handshake to a byte array
		/// </summary>
		internal byte[] ToBytes()
		{
			// construct our handshake
			var handshake = string.Format( this.ClientHandshakeFormat, this.ResourcePath, this.Host, this.Key1, this.Key2 );
			var handshakeBytes = Encoding.UTF8.GetBytes( handshake ).ToList();

			// add in the handshake challenge
			handshakeBytes.AddRange( this.Challenge );

			return handshakeBytes.ToArray();
		}
		#endregion
		
		#region Internal Properties
		/// <summary>
		/// The host name of the server
		/// </summary>
		internal string Host { get; set; }

		/// <summary>
		/// The origin of the client connection
		/// </summary>
		internal string Origin { get; set; }

		/// <summary>
		/// The resource path being used in the connection
		/// </summary>
		internal string ResourcePath { get; set; }

		/// <summary>
		/// A randomly generated key for authenticating the handshake
		/// </summary>
		internal string Key1 { get; set; }

		/// <summary>
		/// A randomly generated key for authenticating the handshake
		/// </summary>
		internal string Key2 { get; set; }

		/// <summary>
		/// A randomly generated key for authenticating the handshake
		/// </summary>
		internal byte[] Challenge { get; set; }
		#endregion

		#region Private Methods
		/// <summary>
		/// Generates a key for the handshake
		/// </summary>
		private string GenerateHandshakeKey()
		{
			var seed = BitConverter.ToInt32( this.GetRandomBytes( 8 ), 0 );
			var random = new Random( seed );
			var randomChars = this.GetAcceptableCharacters();

			var spaces = ( random.NextDouble() * 12 ) + 1;
			var max = 4294967295 / spaces;
			var number = random.NextDouble() * max;
			var product = number * spaces;
			var key = Math.Truncate( product ).ToString();

			var randomInserts = ( random.NextDouble() * 12 ) + 1;
			for ( var i = 0; i < randomInserts; i++ )
			{
				var pos = (int)( random.NextDouble() * ( key.Length + 1 ) );
				var chPos = (int)( random.NextDouble() * randomChars.Length );
				key = key.Insert( pos, randomChars.Substring( chPos, 1 ) );
			}

			var spaceChar = " ";
			for ( var i = 0; i < spaces; i++ )
			{
				var pos = (int)( random.NextDouble() * ( key.Length - 1 ) ) + 1;
				key = key.Insert( pos, spaceChar );
			}

			return key;
		}

		/// <summary>
		/// Gets a list of the characters usable in the handshake key
		/// </summary>
		private string GetAcceptableCharacters()
		{
			var charMap = new char[84];

			for ( int i = 0x0021; i <= 0x002F; i++ )
				charMap[i - 0x0021] = (char)i;

			int existingLength = 0x002F - 0x0021 + 1;

			for ( int i = 0x003A; i <= 0x007E; i++ )
				charMap[i - 0x003A + existingLength] = (char)i;

			return new string( charMap );
		}

		/// <summary>
		/// Generate a random series of bytes
		/// </summary>
		private byte[] GetRandomBytes( int size )
		{
			var random = new System.Security.Cryptography.RNGCryptoServiceProvider();
			byte[] randomBytes = new byte[size];

			random.GetBytes( randomBytes );

			return randomBytes;
		}
		#endregion

		/// <summary>
		/// A format string for construct the client handshake
		/// </summary>
		internal string ClientHandshakeFormat = "GET {0} HTTP/1.1\r\n" + 
												"Upgrade: WebSocket\r\n" + 
												"Host: {1}\r\n" + 
												"Sec-WebSocket-Key1: {2}\r\n" + 
												"Sec-WebSocket-Key2: {3}\r\n" + 
												"Origin: null\r\n" +
												"Connection: Upgrade\r\n" + 
												"\r\n";

		/// <summary>
		/// Extracts the client handshake values from a byte buffer
		/// </summary>
		internal static WebSocketClientHandshake Parse( byte[] clientHandshake, int clientHandshakeSize )
		{
			// convert the handshake to text for regular expressions searching
			var handshake = Encoding.UTF8.GetString( clientHandshake );

			var webSocketClientHandshake = new WebSocketClientHandshake()
			{
				ResourcePath = Regex.Match( handshake, @"GET (?<resourcePath>.*) HTTP/1.1\r\n" ).Groups["resourcePath"].Value,
				Host = Regex.Match( handshake, @"Host: (?<host>.*)\r\n" ).Groups["host"].Value,
				Origin = Regex.Match( handshake, @"Origin: (?<origin>.*)\r\n" ).Groups["origin"].Value,
				Key1 = Regex.Match( handshake, @"Sec-WebSocket-Key1: (?<key>.*)\r\n" ).Groups["key"].Value,
				Key2 = Regex.Match( handshake, @"Sec-WebSocket-Key2: (?<key>.*)\r\n" ).Groups["key"].Value,
				Challenge = clientHandshake.Skip( clientHandshakeSize - 8 ).Take( 8 ).ToArray()
			};

			return webSocketClientHandshake;
		}		
	}
}
